#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
shb_ns3_client_phone.py —— 使用 unified_netgear 的 NetGearLike 封装发送数据
本版更新要点：
  - send_file_via_netgear：loss_rate 从 RTCP 获取；rtt_ms 从底层 PING/PONG 获取；
  - 'tooth' 分支：调用 net.get_slow_prediction() 使用慢模块最新建议的 fec_suggest；
  - 其它逻辑沿用你之前版本；不直接导入 PKT_*，通过 net.is_* 判断类型。

  # MODIFIED（GCC）：
  - 新增 --gcc 选项；为真时：
      * 启用 net.enable_gcc(True)
      * 模拟“编码器”效果：根据目标码率 A（bps）将 frame_data 调整到 A/8/fps 字节
      * 发送分片改用 net.enqueue_frame_pkts(pkts, frame_id) 走 pacer（节奏化发送）

  # 本次补丁（最小变更）：
      * CHANGED: bytes_total 只统计 payload（分片 len(raw) 之和），不再叠加 APP 头
      * CHANGED: origin_bytes 在 GCC 裁剪/填充之后再记录
      * CHANGED: _try_flush() 支持无 ACK 时在 10s 超时后估算 total_ms 并落盘
      * CHANGED: finally 阶段做干净退出：强制 flush、join 接收线程、net.close()
"""

import argparse
import csv
import math
import os
import struct
import threading
import time
import socket
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from collections import deque
import re
from gears.netgear_udp import NetGearUDP as NetGear
# from gears.tooth_fast_module import ToothFastModuleRF
from gears.webrtc_interface import WebRtcPolicy

from ns3.sender import sendFrame, sendFrame_xor  # 分片（可选）

UDP_BIND_ADDR: str = "0.0.0.0"
UDP_BIND_PORT: int = 5559
UDP_DST_ADDR: str = "114.212.86.152" # 127.0.0.1,"114.212.86.152"
UDP_DST_PORT: int = 5558

DEFAULT_MTU: int = 1500
MAX_PAYLOAD: int = 1400
APP_HEADER_BYTES: int = 29
NS3_DEFAULT_BITRATE_MBPS: float = 10.0

_ACK_FMT = "!id"
_ACK_SIZE = struct.calcsize(_ACK_FMT)

_GLOBAL_UDP = None
_UDP_RX_THREAD: Optional[threading.Thread] = None
_UDP_RX_TERMINATE = False

# === 新增：ACK/RTCP 统计（限频打印，避免 I/O 成为瓶颈）===
_ACK_FRAME_RECV = 0       # 帧级 ACK 计数
_PRINT_EVERY = 50         # 每收到多少个再打印一次

_PER_FRAME: Dict[int, Dict[str, Any]] = {}
_PER_LOCK = threading.Lock()

_CLIENT_LOG = Path(f"./client_times_{time.strftime('%Y%m%d_%H%M%S')}.csv")
# === NEW: 统一会话名（供 NetGear 内部输出文件命名使用） ===
_SESSION_NAME: str = f"default_{int(time.time())}"  # 含义：会话名，默认时间戳占位
_MIN_FRAMES_PER_FILE: int = 18000
# === 每个输入文件专属：用于“输出文件 frame_id 从 0 开始”的写盘视角 ===
_CURRENT_FILE_FID_OFFSET: int = 0   # 本输入文件的起始偏移（写出时要减去）
# ---- 全局 GCC 开关（由 --gcc 设置）----
_USE_GCC = False
_USE_VIDUCE = False
def _pf_update(fid: int, **kv: Any) -> None:
    with _PER_LOCK:
        rec = _PER_FRAME.setdefault(int(fid), {})
        rec.update(kv)

def _try_flush(fid: int) -> None:
    """
    CHANGED: 支持“无 ACK 超时（10s）”的估算落盘。
    规则：
      - 必须具备：bytes_total、t_send_start、t_send_end
      - 若 total_ms 已有：直接落盘
      - 若 total_ms 缺失且 (now - t_send_start) >= 10s：
            估算 total_ms = send_ser_measured_ms + (ping_rtt_ms 若可得，否则 0)
            然后落盘，避免卡死等待 ACK
    """
    with _PER_LOCK:
        rec = _PER_FRAME.get(int(fid), {})
        if not rec:
            return
        origin_bytes= rec.get("origin_bytes")
        bytes_total = rec.get("bytes_total")
        total_ms    = rec.get("total_ms")
        t0          = rec.get("t_send_start")
        ping_rtt_ms = rec.get("ping_rtt_ms", None)
        bytes_sum = rec.get("bytes_sum")

        # 基本条件：需要 payload 统计与起止时间
        if (bytes_total is None) or (t0 is None):
            return

        # 若缺少 total_ms，则在 3s 后做兜底估算（无 ACK）
        if total_ms is None:
            now = time.time()
            if (now - float(t0)) < 3.0:
                return  # 未超时就继续等 ACK
            # 丢包，直接写NAN
            rec["total_ms"] = float("nan")  # 写回，后续统一按“有 total_ms”的路径落盘

        with _CLIENT_LOG.open("a", newline="") as f:
            w = csv.writer(f)
            fid_for_csv = int(fid) - int(_CURRENT_FILE_FID_OFFSET)
            w.writerow([
                int(fid_for_csv),
                int(origin_bytes) if origin_bytes is not None else 0,
                int(bytes_total),
                int(bytes_sum),
                ("NA" if math.isnan(rec["total_ms"]) else round(float(rec["total_ms"]))),
                ("NA" if ping_rtt_ms is None else round(float(ping_rtt_ms), 3)),
            ])
        del _PER_FRAME[int(fid)]

def _force_flush_all_pending_as_na() -> None:
    """
    在等待 10s 之后调用：将仍未落盘（无 total_ms）的帧强制标记为 NA 并写出。
    随后清空 _PER_FRAME，确保不会与下一个输入文件的写盘相互污染。
    """
    with _PER_LOCK:
        pending_ids = list(_PER_FRAME.keys())

    for fid in pending_ids:
        with _PER_LOCK:
            rec = _PER_FRAME.get(int(fid), {})
            if rec and ("total_ms" not in rec or rec["total_ms"] is None):
                # 标记为 NA，交给 _try_flush 正常写盘
                rec["total_ms"] = float("nan")
                _PER_FRAME[int(fid)] = rec
        _try_flush(int(fid))

    # 清空当前文件的在途状态（避免跨文件混写）
    with _PER_LOCK:
        _PER_FRAME.clear()

# -------------------- UDP 入口 --------------------
def ensure_global_udp() -> Any:
    """创建/复用全局 UDP 会话（你的 NetGearLike 封装）。"""
    global _GLOBAL_UDP, _UDP_RX_THREAD, _UDP_RX_TERMINATE
    if _GLOBAL_UDP is not None:
        return _GLOBAL_UDP
    if NetGear is None:
        raise RuntimeError("未找到 NetGear，请确认 vidgear 环境。")

    _GLOBAL_UDP = NetGear(
        address=UDP_DST_ADDR,
        port=UDP_DST_PORT,
        protocol="udp",
        receive_mode=False,
        local_bind_port = UDP_BIND_PORT,
        logging=True,
        mtu=DEFAULT_MTU,
        send_buffer_size=32 * 1024 * 1024,
        recv_buffer_size=32 * 1024 * 1024,
        queue_maxlen=65536,
        slow_log_dir="./tmp/logs_sender",            # 发送端日志
        slow_weights_path="./tmp/slow_module_weights.json",
        rtcp_interval_ms=100,
        ts_name=_SESSION_NAME, 
    )
    _GLOBAL_UDP.set_target_bps(NS3_DEFAULT_BITRATE_MBPS * 1e6)  # 初始目标码率
    # 指定对端
    try:
        setattr(_GLOBAL_UDP, "_peer_addr", (UDP_DST_ADDR, UDP_DST_PORT))
    except Exception:
        print("set peer addr fail!")
    try:
        if hasattr(_GLOBAL_UDP, "set_peer"):
            _GLOBAL_UDP.set_peer(UDP_DST_ADDR, UDP_DST_PORT)  # type: ignore
    except Exception:
        print("get peer addr fail!")

    # NEW: 若支持，启用 GCC（仅在 --gcc 时）
    global _USE_GCC
    try:
        if _USE_GCC and hasattr(_GLOBAL_UDP, "enable_gcc"):
            _GLOBAL_UDP.enable_gcc(True)
    except Exception:
        pass

    # 后台接收线程（专用、轻量、仅做分发与状态更新）
    _UDP_RX_TERMINATE = False
    _UDP_RX_THREAD = threading.Thread(target=_client_rx_loop, args=(_GLOBAL_UDP,), daemon=True)
    _UDP_RX_THREAD.start()
    return _GLOBAL_UDP

def get_last_rtt_ms() -> float:
    """返回最近 RTT（ms）：优先 median；退化用 ema。"""
    try:
        net = ensure_global_udp()
        st = net.get_rtt_stats() or {}
        v = st.get("median_ms") or st.get("ema_ms")
        return float(v) if v is not None else float("nan")
    except Exception:
        return float("nan")

# -------------------- RTCP 拉取（用于发送函数中取最近一条） --------------------
def _pump_all_rtcp_and_get_latest(net: Any) -> Dict[str, Any]:
    """抽干 RTCP 队列并返回最近一条；若无则返回空 dict。"""
    last = None
    while True:
        try:
            rep = net.get_rtcp_report()
        except Exception:
            rep = None
        if rep is None:
            break
        last = rep
    return {} if last is None else dict(last)

# -------------------- WebRTC 策略 --------------------
def compute_webrtc_fec_rate_strict(loss_rate: float, rtt_ms: float, fps: float, bitrate_mbps: float) -> float:
    if WebRtcPolicy is None:
        return 0.0
    try:
        return float(WebRtcPolicy.compute_fec(loss_rate, rtt_ms, fps, bitrate_mbps))
    except Exception:
        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        return 0.0

_RTT_frame  = 0
# -------------------- 轻量事件处理：ACK/RTCP/RES/PONG --------------------
def _on_ack_frame(net: Any, data: bytes) -> None:
    """帧级 ACK：仅做轻处理与统计，禁止重 I/O。"""
    global _ACK_FRAME_RECV
    global _RTT_frame
    if len(data) < _ACK_SIZE:
        return
    try:
        fid, server_send_ts = struct.unpack(_ACK_FMT, data[:_ACK_SIZE])
        now = time.time()
        _pf_update(int(fid), t_ack_recv=float(now))
        rec = _PER_FRAME.get(int(fid), {})
        t0 = rec.get("t_send_start")

        if t0 is not None:
            total_ms = max(0.0, (now - float(t0)) * 1000.0)
            _RTT_frame = total_ms
            _pf_update(int(fid), total_ms=float(total_ms))
        # RTT（轻量获取）
        try:
            st = net.get_rtt_stats() or {}
            rtt_ms = st.get("median_ms") or st.get("ema_ms")
            if rtt_ms is not None:
                rtt_ms = float(rtt_ms)
                _pf_update(int(fid), ping_rtt_ms=rtt_ms)
        except Exception:
            pass
        _try_flush(int(fid))

        _ACK_FRAME_RECV += 1
        if (_ACK_FRAME_RECV % 10) == 0:
            # 限频打印，避免 I/O 阻塞
            print(f"[Client] ACK_FRAME received = {_ACK_FRAME_RECV}")
    except Exception:
        pass

def _on_res(net: Any, data: bytes) -> None:
    """RES：按需解析（此处保持轻处理，避免 I/O）"""
    # if len(data) >= 8 and data[:4] == b'RSID':
    #     _, fid = struct.unpack('!4sI', data[:8])  # 稳定获得 frame_id
    #     now = time.time()
    #     _pf_update(int(fid), t_res_recv=float(now))
    #     rec = _PER_FRAME.get(int(fid), {})
    #     t0 = rec.get('t_send_start')
    #     if t0 is not None:
    #         res_ms = max(0.0, (now - float(t0)) * 1000.0)
    #         _pf_update(int(fid), res_elapsed_ms=float(res_ms))
    #     _try_flush(int(fid))
    return


# -------------------- 专用接收线程（轻量抽水） --------------------
def _client_rx_loop(net: Any) -> None:
    """
    发送端的专用接收线程：持续抽干 net.recv()，避免 ACK/RTCP 堵在缓冲里。
    - 严禁在这里做重计算/磁盘 I/O
    - 仅做轻解析与状态更新
    """
    global _UDP_RX_TERMINATE
    while not _UDP_RX_TERMINATE:
        try:
            item = net.recv()
        except Exception:
            item = None
        if not item:
            time.sleep(0.0005)
            continue
        pkt_type = int(item.get("pkt_type", -1))
        data: bytes = item.get("data", b"")

        # 分发（保持轻处理）
        if net.is_ack_frame(pkt_type):
            _on_ack_frame(net, data)
            continue

        if net.is_res(pkt_type):
            _on_res(net, data)
            continue

# -------------------- 核心：发送一帧（按要求修改） --------------------
def send_file_via_netgear(frame_data: bytes, frame_id: int, fps: float, fec_policy: str) -> None:
    """
    修改点：
      - loss_rate：从 RTCP 获取（抽干队列，取最近一条；若无则 0）
      - rtt_ms   ：通过 net.get_rtt_stats()（PING/PONG 估计）
      - 'tooth' 分支：从 net.get_slow_prediction() 取 fec_suggest
      - NEW: 若 _USE_GCC 为真，先模拟“编码器”输出码率 = A，然后再参与分片/FEC；分片下发改为 enqueue_frame_pkts()
      - CHANGED: origin_bytes 在 GCC 裁剪/填充之后再记录
      - CHANGED: bytes_total 仅统计 payload（各分片 len(raw) 之和）
    """
    net = ensure_global_udp()
    if len(frame_data)==0:
        _pf_update(int(frame_id), origin_bytes=0)
        _pf_update(int(frame_id), bytes_total=0)
        _pf_update(int(frame_id), bytes_sum=0)
        _pf_update(int(frame_id), total_ms=0)
        _pf_update(int(frame_id), t_send_start=time.time())
        st = net.get_rtt_stats() or {}
        rtt_ms = st.get("median_ms") or st.get("ema_ms")
        if rtt_ms is not None:
            rtt_ms = float(rtt_ms)
            _pf_update(int(frame_id), ping_rtt_ms=rtt_ms)
        else:
            _pf_update(int(frame_id), ping_rtt_ms=None)
        _try_flush(int(frame_id))
        return

    # —— RTCP 最近一条 —— 
    last_rtcp = _pump_all_rtcp_and_get_latest(net)
    if last_rtcp:
        loss_rate = float(last_rtcp.get("lr", 0.0))
        la = float(last_rtcp.get("la", 0.0))
    else:
        loss_rate = 0.0
        la = 0.0

    # —— RTT（由 PING/PONG 得到；你的 RTCP 未携带 RTT 字段）——
    st = net.get_rtt_stats() or {}
    # TODO: 测一下哪个更好
    rtt_ms = st.get("median_ms") or st.get("ema_ms") or 0.0
    rtt_ms = float(rtt_ms)
    # rtt_ms = _RTT_frame

    _pf_update(int(frame_id), origin_bytes=int(len(frame_data)))
    # ===== NEW: GCC 模拟编码器 =====
    global _USE_GCC
    if _USE_GCC and hasattr(net, "get_gcc_target_bps"):
        try:
            A_bps = float(net.get_gcc_target_bps())
            if not math.isfinite(A_bps) or A_bps <= 0:
                A_bps = NS3_DEFAULT_BITRATE_MBPS * 1e6
        except Exception:
            A_bps = NS3_DEFAULT_BITRATE_MBPS * 1e6
        # 按帧目标大小：A/8/fps （字节）
        target_bytes = int(max(1, A_bps / 8.0 / max(1e-6, float(fps))))
        target_bytes = max(target_bytes, round(len(frame_data)*0.9))
        # 将 frame_data 调整到 target_bytes：过大则截断；过小则零填充
        if len(frame_data) > target_bytes:
            frame_data = frame_data[:target_bytes]
        # 同时将 A 写回底层（可选）
        try:
            if hasattr(net, "set_target_bps"):
                net.set_target_bps(A_bps)
        except Exception:
            pass

    # CHANGED: 在可能的 GCC 裁剪/填充之后再记录 origin_bytes，保证口径一致
    _pf_update(int(frame_id), bytes_total=int(len(frame_data)))

    bitrate_mbps = NS3_DEFAULT_BITRATE_MBPS
    if _USE_GCC:
        media_bps = max(1e5, float(len(frame_data)) * 8.0 * float(fps))
        media_bps = media_bps * 1.2  # 预留 20% 余量
        bitrate_mbps = media_bps / 1e6
        bitrate_mbps = float(net.get_gcc_target_bps() / 1e6)

    

    # —— 选择 FEC 策略 —— 
    if fec_policy == "no-fec":
        fec_rate = 0.0

    elif fec_policy == "webrtc":
        fec_rate = compute_webrtc_fec_rate_strict(loss_rate=loss_rate, rtt_ms=rtt_ms,
                                                  fps=float(fps), bitrate_mbps=bitrate_mbps)

    elif fec_policy == "hairpin":
        # ddl_ms = 100.0
        # k = 500.0
        # webrtc_fec_rate = compute_webrtc_fec_rate_strict(loss_rate=loss_rate, rtt_ms=rtt_ms,
        #                                                  fps=float(fps), bitrate_mbps=bitrate_mbps)
        # fec_rate = max(0.0, min(3.0, webrtc_fec_rate * k * (rtt_ms / max(1.0, (ddl_ms - rtt_ms)))))
        webrtc_fec_rate = compute_webrtc_fec_rate_strict(loss_rate=loss_rate, rtt_ms=rtt_ms,
                                                         fps=float(fps), bitrate_mbps=bitrate_mbps)
        data_size = len(frame_data)
        n = data_size // MAX_PAYLOAD + int(data_size % MAX_PAYLOAD != 0)
        fec_num = round(n*webrtc_fec_rate)+3
        fec_rate = fec_num / n
    elif fec_policy == "tooth":
        webrtc_fec_rate = compute_webrtc_fec_rate_strict(loss_rate=loss_rate, rtt_ms=rtt_ms,
                                                         fps=float(fps), bitrate_mbps=bitrate_mbps)
        data_size = len(frame_data)
        n = data_size // MAX_PAYLOAD + int(data_size % MAX_PAYLOAD != 0)
        fec_num = round(n*webrtc_fec_rate)+2
        fec_rate = fec_num / n

    # elif fec_policy == "tooth":
    #     slow_pred = net.get_slow_prediction()  # {'lr_f','la_f'}
    #     lr_f = float(slow_pred.get("lr_f", 0.0))
    #     la_f = float(slow_pred.get("la_f", 0.0))
    #     fl_pkts = max(1, math.ceil(len(frame_data) / MAX_PAYLOAD))
    #     try:
    #         FAST_MODEL_PATH = "./fast_module.joblib"
    #         _FAST = getattr(send_file_via_netgear, "_FAST_MODEL", None)
    #         if _FAST is None:
    #             _FAST = ToothFastModuleRF.load(FAST_MODEL_PATH)
    #             setattr(send_file_via_netgear, "_FAST_MODEL", _FAST)
    #         fec_rate = float(_FAST.predict(lr_f=lr_f, la_f=la_f, fl_pkts=fl_pkts))
    #         if not (fec_rate == fec_rate) or fec_rate in (float("inf"), float("-inf")) or fec_rate < 0.0:
    #             raise ValueError(f"fast-module predicted invalid fec={fec_rate}")
    #         fec_rate = max(0.0, min(3.0, fec_rate))
    #         print(f"[TOOTH] lr_f={lr_f:.4f} la_f={la_f:.3f} fl_pkts={fl_pkts} -> fec={fec_rate:.3f}")
    #     except Exception as e:
    #         fec_rate = compute_webrtc_fec_rate_strict(
    #             loss_rate=lr_f,
    #             rtt_ms=rtt_ms,
    #             fps=float(fps),
    #             bitrate_mbps=bitrate_mbps
    #         )
    #         fec_rate = max(0.0, min(3.0, float(fec_rate)))
    else:
        print("Error! No fec policy match!")
        fec_rate = 0.0

    # —— 分片与发送 —— 
    if _USE_VIDUCE:
        pkts: List[bytes] = sendFrame_xor(frame_data, loss_rate=loss_rate, rtt_ms=rtt_ms,
                                          fec_rate=fec_rate, max_pay_load=MAX_PAYLOAD, frame_id = frame_id)
    else:
        pkts: List[bytes] = sendFrame(frame_data, loss_rate=loss_rate, rtt_ms=rtt_ms,
                                        fec_rate=fec_rate, max_pay_load=MAX_PAYLOAD, frame_id = frame_id)

    t0 = time.time()
    _pf_update(int(frame_id), t_send_start=float(t0))

    bytes_sum = 0  # CHANGED: 仅累计 payload 字节（各分片 len(raw) 之和）

    # === TODO: 不启用 GCC，使用 pacer 统一节奏化发送；沿用原逻辑逐包 net.send ===
    if _USE_GCC and hasattr(net, "enqueue_frame_pkts"):
        net.enqueue_frame_pkts(pkts, frame_id=int(frame_id))
        # CHANGED: 只统计 payload
        for raw in pkts:
            bytes_sum += (len(raw)+15)
    else:
        sleep_time = min(0.0001, 0.025 / len(pkts))
        for i, raw in enumerate(pkts):
            net.send(raw)  # 默认 DATA
            bytes_sum += (len(raw)+15)
            # time.sleep(0)

    _pf_update(int(frame_id), bytes_sum=bytes_sum)

    _try_flush(int(frame_id))
    if frame_id % 50 == 0:
        print(f"[Sender] frame_id={frame_id} bytes={bytes_sum} "
              f"lr(rtpc)={loss_rate:.4f} rtt(ping)={rtt_ms:.1f} rtt(ack)={_RTT_frame:.1f} fec={fec_rate:.3f} cur bitrate: {bitrate_mbps} Mbps")

# -------------------- 其它：回放/主函数（与原版一致，新增 --gcc） --------------------
def _fallback_sendFrame(frame_data: bytes, max_pay_load: int) -> List[bytes]:
    pkts: List[bytes] = []
    n = len(frame_data); off = 0
    while off < n:
        pkts.append(frame_data[off: off+max_pay_load])
        off += max_pay_load
    return pkts

# NEW: 从文件名中提取序号（最后一个下划线后的数字）。例如：c1_X265NonB_downsampling_00000.txt -> 0
_NUM_RE = re.compile(r".*_([0-9]+)$")

def _extract_seq_num(path: str) -> int:
    base = os.path.splitext(os.path.basename(path))[0]
    m = _NUM_RE.match(base)
    return int(m.group(1)) if m else -1

def _load_sizes_from_file(path: str) -> List[Tuple[int, int]]:
    """
    输入文件格式要求：
      - 每行至少两列：fid, size
      - 允许逗号分隔或空白分隔
      - 忽略空行与无效行
    返回：[(fid, size), ...]
    """
    pairs: List[Tuple[int, int]] = []
    if not os.path.exists(path):
        raise FileNotFoundError(f"metrics 文件不存在: {path}")

    with open(path, "r", newline="") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            # 同时支持逗号与空白
            parts = [p for p in re.split(r"[,\s]+", line) if p]
            if len(parts) < 2:
                continue
            try:
                fid = int(parts[0])
            except Exception:
                fid = len(pairs)
            try:
                size = int(parts[1])
            except Exception:
                continue
            if size < 0:
                continue
            pairs.append((fid, size))

    if not pairs:
        raise RuntimeError(f"未能在 {path} 解析出任何有效的 (fid,size)")
    return pairs

def _random_payload(nbytes: int) -> bytes:
    return os.urandom(int(nbytes))

def replay_from_metrics(csv_path: str, interval_ms: float = 30.0, fec_policy: str = "webrtc",
                        fid_offset: int = 0, total_frames: Optional[int] = None) -> None:
    """
    回放 csv_path 中的 (fid,size) 记录；当 total_frames > 源文件帧数时，按模循环 size，直到凑够 total_frames 帧。
    发送用的 frame_id = fid_offset + i（i 从 0 连续递增），确保全局唯一；写盘侧会再减去“本文件偏移”，从 0 开始。
    """
    net = ensure_global_udp()
    sizes = _load_sizes_from_file(csv_path)
    n = len(sizes)
    if n <= 0:
        return

    # 目标总帧数：未指定则取源长度；小于最小阈值时，循环补齐
    if total_frames is None:
        total_frames = n
    if total_frames < _MIN_FRAMES_PER_FILE:
        total_frames = _MIN_FRAMES_PER_FILE

    fps = 1000.0 / float(interval_ms)
    next_ts = time.time()

    # i 是“本文件内的发送序号”，不会用源文件自带 fid
    for i in range(int(total_frames)):
        # 循环取 size：当源不足时按模复用
        _, size = sizes[i % n]

        # ### 可选调试：固定 size
        # size = 1000

        payload = _random_payload(size)
        # 发送 frame_id：全局唯一 = 外层传入的文件偏移 + 本文件内的 i
        send_file_via_netgear(payload, frame_id=int(fid_offset) + int(i), fps=fps, fec_policy=fec_policy)

        # 节奏控制
        next_ts += interval_ms / 1000.0
        sleep_s = next_ts - time.time()
        if sleep_s > 0:
            time.sleep(sleep_s)


def receive_weights(listen_ip: str, listen_port: int, save_dir: str = "/data/shb/test/receive"):
    """
    持续通过 TCP 接收权重文件，并打印每次文件接收的开始时间、结束时间与耗时。
    """
    os.makedirs(save_dir, exist_ok=True)

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
        server_socket.bind((listen_ip, listen_port))
        server_socket.listen(5)
        print(f"📡 持续监听中 {listen_ip}:{listen_port} ...")

        while True:
            conn, addr = server_socket.accept()
            print(f"\n🔗 已连接来自: {addr}")

            try:
                # --- 接收文件头部信息 ---
                header = b""
                while not header.endswith(b"\n"):
                    data = conn.recv(1)
                    if not data:
                        break
                    header += data

                if not header:
                    print("⚠️ 未收到任何数据，可能客户端异常断开。")
                    conn.close()
                    continue

                file_name, file_size = header.decode().strip().split("|")
                file_size = int(file_size)
                save_path = os.path.join(save_dir, file_name)

                # --- 记录开始时间 ---
                start_time = time.time()
                start_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(start_time))
                print(f"🕒 开始接收文件: {file_name}  ({file_size} 字节)")
                print(f"⏱️ 接收开始时间: {start_str}")

                # --- 接收文件内容 ---
                received_bytes = 0
                with open(save_path, "wb") as f:
                    while received_bytes < file_size:
                        chunk = conn.recv(8388608)
                        if not chunk:
                            break
                        f.write(chunk)
                        received_bytes += len(chunk)

                # --- 记录结束时间与耗时 ---
                end_time = time.time()
                end_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(end_time))
                duration = end_time - start_time

                print(f"✅ 文件接收完成: {save_path}")
                print(f"📥 接收结束时间: {end_str}")
                print(f"⏳ 接收总耗时: {duration:.2f} 秒\n")

            except Exception as e:
                print(f"⚠️ 接收出错: {e}")
            finally:
                conn.close()


def receive_weights_udp(listen_ip: str, listen_port: int, save_dir: str = "/data/shb/test/receive"):
    """
    通过 UDP 接收文件。
    文件需由 send_weights_udp() 发送。
    支持自动识别 HEADER / DATA / END 包并保存文件。
    """
    os.makedirs(save_dir, exist_ok=True)

    # 创建 UDP socket
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_socket.bind((listen_ip, listen_port))

    print(f"📡 UDP 接收端已启动，监听 {listen_ip}:{listen_port} ...")

    current_file = None
    file_name = None
    expected_size = 0
    received_bytes = 0
    start_time = None

    # 存储数据包顺序（避免乱序）
    data_buffer = {}

    while True:
        data, addr = server_socket.recvfrom(65535)  # 单个UDP包最大长度
        message = data.decode(errors="ignore")

        # --- 处理文件头 ---
        if message.startswith("HEADER|"):
            parts = message.strip().split("|")
            if len(parts) != 3:
                print("⚠️ 文件头格式错误，跳过。")
                continue

            _, file_name, size_str = parts
            expected_size = int(size_str)
            save_path = os.path.join(save_dir, file_name)
            current_file = open(save_path, "wb")
            received_bytes = 0
            data_buffer.clear()
            start_time = time.time()

            print(f"\n🕒 开始接收文件: {file_name} ({expected_size} 字节)")
            print(f"来自: {addr}")

        # --- 处理数据块 ---
        elif message.startswith("DATA|"):
            # 格式: DATA|seq|<二进制数据>
            header_end = message.find("|", 5)  # 第二个分隔符位置
            seq_str = message[5:header_end]
            seq = int(seq_str)
            content = data[header_end+1:]

            if seq not in data_buffer:
                data_buffer[seq] = content
                received_bytes += len(content)

                # 实时写入（简单实现：按顺序写）
                current_file.write(content)

        # --- 处理结束标志 ---
        elif message.strip() == "END":
            if current_file:
                current_file.close()
                end_time = time.time()
                duration = end_time - start_time

                print(f"✅ 文件接收完成: {file_name}")
                print(f"📦 实际接收: {received_bytes} 字节 / 期望: {expected_size} 字节")
                print(f"⏳ 用时: {duration:.2f} 秒, 平均速率: {received_bytes / duration / 1024:.2f} KB/s\n")

                current_file = None
                file_name = None
                expected_size = 0
                received_bytes = 0
                data_buffer.clear()

        else:
            print(f"⚠️ 未识别的数据包: {message[:50]} ...")

def main() -> None:
    global _UDP_RX_TERMINATE, _UDP_RX_THREAD, _GLOBAL_UDP
    global _CLIENT_LOG, _SESSION_NAME, _CURRENT_FILE_FID_OFFSET
    parser = argparse.ArgumentParser()
    parser.add_argument("--metrics_folder", type=str, default="/home/wxk/workspace/nsdi/Viduce/net/vidgear/frame_lengths",
                        help="包含若干 .txt / .csv 的文件夹，每个文件含 (fid,size)")
    parser.add_argument("--output_dir", type=str, default="/home/wxk/workspace/nsdi/Viduce/net/vidgear/tmp",
                        help="输出目录，逐文件生成 net_{fec_policy}_{input_name}.csv")
    parser.add_argument("--interval_ms", type=float, default=30.0,
                        help="两帧间隔（毫秒）")
    parser.add_argument("--fec_policy", type=str, default="webrtc",
                        help="no-fec, webrtc, hairpin, tooth")
    # ===== NEW: GCC 开关 =====
    parser.add_argument("--gcc", action="store_true",
                        help="启用 GCC：模拟编码器 + pacer 节奏化发送 + REMB/As 融合")
    parser.add_argument("--viduce", action="store_true",
                        help="启用 Viduce")
    args = parser.parse_args()
    # 设置全局 GCC flag
    args.gcc = True
    global _USE_GCC
    _USE_GCC = True#bool(args.gcc)
    global _USE_VIDUCE
    _USE_VIDUCE = bool(args.viduce)
    # 确保输出目录存在
    os.makedirs(args.output_dir, exist_ok=True)


    # 递归搜集所有名为 frame_length.csv 的文件（多级目录）
    folder = Path(args.metrics_folder)
    files = [p for p in folder.rglob("frame_length.csv")]
    files.sort()  # 以路径字典序排序，保证稳定处理顺序

    if not files:
        raise RuntimeError(f"未在 {folder} 及其子目录找到任何 frame_length.csv")

    # 全局累计偏移：保证跨文件的 frame_id 不重复
    cumulative_fid_offset = 0
    try:
        print(f"------------------ {args.fec_policy}, viduce: {_USE_VIDUCE} ------------------")
        for in_path in files:
            in_path = Path(in_path)
            # 计算相对路径，构造镜像输出目录
            rel_parent = in_path.parent.relative_to(folder)  # 输入根目录下的相对父路径
            out_dir = Path(args.output_dir) / rel_parent
            out_dir.mkdir(parents=True, exist_ok=True)

            # 针对“每个输入文件”设置独立的输出 CSV
            # 目标：c1\30\car_camera\libvpx\DDS -> "c1_30_car_camera_libvpx_DDS"
            rel_parts = list(rel_parent.parts)  # e.g. ["c1","30","car_camera","libvpx","DDS"]
            # 清洗每一段的非法字符（只保留字母数字和短横线，把其他都变成下划线）
            rel_parts = [re.sub(r'[^0-9A-Za-z\-]+', '_', p).strip('_') for p in rel_parts]
            rel_name = "_".join([p for p in rel_parts if p]) or "root"

            for i in range(144):
                if _USE_VIDUCE:
                    out_csv = out_dir / f"net_viduce+{args.fec_policy}_{rel_name}_{i}.csv"
                else:
                    out_csv = out_dir / f"net_{args.fec_policy}_{rel_name}_{i}.csv"


                # === 关键：设置“当前文件”的写盘视角偏移 ===
                _CLIENT_LOG = out_csv
                _SESSION_NAME = f"{args.fec_policy}_{rel_name}_{int(time.time())}"
                _CURRENT_FILE_FID_OFFSET = int(cumulative_fid_offset)

                # 输出 CSV 头
                with _CLIENT_LOG.open("w", newline="") as f:
                    w = csv.writer(f)
                    w.writerow([
                        "frame_id","origin_bytes","finetuned_bytes","sent_bytes","total_ms",
                        "ping_rtt_ms"
                    ])

                # 读取当前输入文件，得到帧尺寸列表；确定“目标总帧数”（不足则循环补齐到 _MIN_FRAMES_PER_FILE）
                sizes = _load_sizes_from_file(str(in_path))
                num_frames_src = len(sizes)
                num_frames_this_file = max(num_frames_src, _MIN_FRAMES_PER_FILE)

                # 回放当前文件：发送 frame_id = _CURRENT_FILE_FID_OFFSET + i（i=0..num_frames_this_file-1）
                replay_from_metrics(
                    csv_path=str(in_path),
                    interval_ms=float(args.interval_ms),
                    fec_policy=args.fec_policy,
                    fid_offset=int(_CURRENT_FILE_FID_OFFSET),
                    total_frames=int(num_frames_this_file)
                )

                # 文件级 drain：等待 10s，尽可能让 ACK 抵达
                time.sleep(10.0)

                # 强制把剩余未 ACK 的帧按 NA 写出，并清空状态（避免跨文件污染）
                _force_flush_all_pending_as_na()

                # 滚动全局偏移 = 已发送的总帧数
                cumulative_fid_offset += int(num_frames_this_file)
    except KeyboardInterrupt:
        raise SystemExit(130)
    finally:
        # 终止接收线程并 join
        _UDP_RX_TERMINATE = True
        try:
            if _UDP_RX_THREAD is not None:
                _UDP_RX_THREAD.join(timeout=1.0)  # NEW: 等待接收线程退出
        except Exception:
            pass

        # 关闭底层 UDP（会收尾 I/O/RTCP/pacer 等线程，并关闭 socket）
        try:
            if _GLOBAL_UDP is not None and hasattr(_GLOBAL_UDP, "close"):
                _GLOBAL_UDP.close()
        except Exception:
            pass

def start_listener():
    # 启动接收文件的监听进程（运行在独立线程中）
    receive_weights("0.0.0.0", 5555)
    # receive_weights_udp("0.0.0.0", 5555)

if __name__ == "__main__":
    # 创建并启动一个独立的监听线程
    listener_thread = threading.Thread(target=start_listener, daemon=True)
    listener_thread.start()
    main()
    listener_thread.join()
    
# if __name__ == "__main__":
#     main()
